import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import glob
import os
from mpl_toolkits.mplot3d import Axes3D
import imageio.v2 as imageio
import scipy.special as scp  # 添加scipy导入
from scipy.signal import argrelextrema  # 新增导入

class WaveguideVisualizer:
    def __init__(self, frequency):
        self.frequency = frequency
        # 波导参数（与C++代码保持一致）
        self.radius = 0.5  # 波导半径(m)
        self.length = 1.0  # 轴向长度(m)
        self.Nr = 100      # 径向网格数
        self.Nz = 200      # 轴向网格数
        
        # 物理参数（与C++代码保持一致）
        self.c = 299792458.0  # 光速(m/s)
        self.dr = self.radius / (self.Nr - 1)
        self.dz = self.length / (self.Nz - 1)
        self.dt = 1 / (2 * self.c * np.sqrt(1 / (self.dr**2) + 1 / (self.dz**2)))  # 时间步长(s)
        
        # 新增解析解参数
        self.omega = 2 * np.pi * self.frequency
        first_zero = scp.jn_zeros(1, 1)[0]  # J1的第一个零点
        self.kr = first_zero / self.radius
        self.kz = np.sqrt((self.omega / self.c)**2 - self.kr**2)
        self.A0 = 1.0  # 激励幅值
        
        # 创建网格
        self.r = np.linspace(0, self.radius, self.Nr)
        self.z = np.linspace(0, self.length, self.Nz)
        self.R, self.Z = np.meshgrid(self.r, self.z, indexing='ij')
        
        # 创建输出目录和临时目录
        self.output_dir = 'waveguide_plots'
        self.temp_dir = 'temp'
        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs(self.temp_dir, exist_ok=True)

    def load_field(self, filename):
        """加载场数据"""
        return np.loadtxt(filename, delimiter=',')

    def plot_heatmap(self, field, step):
        """绘制热图"""
        plt.figure(figsize=(10, 6))
        plt.pcolormesh(self.Z, self.R, field, shading='auto', cmap='RdBu_r')
        plt.colorbar(label='Electric Field Intensity')
        plt.xlabel('z (m)')
        plt.ylabel('r (m)')
        time_cycles = step * self.dt * self.frequency  # 改为周期单位
        plt.title(f'Waveguide Electric Field Distribution (t = {time_cycles:.2f} periods)')  # 修改标题
        plt.savefig(os.path.join(self.output_dir, 'waveguide_heatmap.png'), dpi=300, bbox_inches='tight')
        plt.close()

    def plot_3d_surface(self, field, step):
        """绘制3D表面图"""
        fig = plt.figure(figsize=(12, 8))
        ax = fig.add_subplot(111, projection='3d')
        surf = ax.plot_surface(self.Z, self.R, field, cmap=cm.coolwarm,
                             linewidth=0, antialiased=True)
        fig.colorbar(surf, label='Electric Field Intensity')
        ax.set_xlabel('z (m)')
        ax.set_ylabel('r (m)')
        ax.set_zlabel('Field Amplitude')
        time_cycles = step * self.dt * self.frequency  # 改为周期单位
        plt.title(f'3D Waveguide Field Distribution (t = {time_cycles:.2f} periods)')  # 修改标题
        plt.savefig(os.path.join(self.output_dir, 'waveguide_3d_surface.png'), dpi=300, bbox_inches='tight')
        plt.close()

    def create_animation(self):
        """创建动画"""
        print("Creating animation...")
        filenames = sorted(glob.glob('data/field_*.csv'))
        
        # 创建帧列表
        frames = []
        
        # 设置固定的图像尺寸和DPI
        fig_width, fig_height = 10, 6
        dpi = 100
        
        for i, filename in enumerate(filenames):
            step = int(filename.split('_')[1].split('.')[0])
            field = self.load_field(filename)
            
            # 创建新的图形，确保尺寸一致
            fig, ax = plt.subplots(1, 2, figsize=(fig_width * 2, fig_height), dpi=dpi)
            plt.sca(ax[0])
            plt.pcolormesh(self.Z, self.R, field, shading='auto', cmap='RdBu_r')
            plt.colorbar(label='Electric Field Intensity')
            plt.xlabel('z (m)')
            plt.ylabel('r (m)')
            time_cycles = step * self.dt * self.frequency
            plt.title(f'Electric Field Distribution (t = {time_cycles:.2f} periods)')
            
            # 添加等离子体分界线
            plasma_boundary = self.length / 4
            ax[0].axvline(x=plasma_boundary, color='red', linestyle='--', linewidth=2, label='Plasma Boundary')
            ax[0].legend(loc='upper right')
            
            # 绘制电流分布
            current_file = filename.replace('field_', 'current_')
            if os.path.exists(current_file):
                current = self.load_field(current_file)
                plt.sca(ax[1])
                plt.pcolormesh(self.Z, self.R, current, shading='auto', cmap='viridis')
                plt.colorbar(label='Current Density')
                plt.xlabel('z (m)')
                plt.ylabel('r (m)')
                plt.title(f'Current Distribution (t = {time_cycles:.2f} periods)')
                
                # 添加等离子体分界线
                ax[1].axvline(x=plasma_boundary, color='red', linestyle='--', linewidth=2, label='Plasma Boundary')
                ax[1].legend(loc='upper right')
            
            # 固定坐标轴范围
            for axis in ax:
                axis.set_xlim(0, self.length)
                axis.set_ylim(0, self.radius)
            
            # 保存为临时文件
            temp_filename = os.path.join(self.temp_dir, f'temp_frame_{i:03d}.png')
            plt.savefig(temp_filename, dpi=dpi)
            plt.close()
            
            # 读取图像并添加到帧列表
            img = imageio.imread(temp_filename)
            frames.append(img)
        
        # 保存动画
        output_path = os.path.join(self.output_dir, 'waveguide_animation.gif')
        try:
            imageio.mimsave(output_path, frames, duration=0.2)  # 每帧持续0.2秒
            print(f"动画已保存到 {output_path}")
        except ValueError as e:
            print(f"保存GIF时出错: {e}")

    def plot_radial_slice(self, field, current, step):
        """绘制并保存径向截面图"""
        r_target = self.radius / 2  # 在半径中间位置取截面
        r_index = np.argmin(np.abs(self.r - r_target))
        electric_field_z = field[r_index, :]
        current_z = current[r_index, :]
        
        # 计算解析解
        t = step * self.dt
        analytical_amplitude = scp.j1(self.kr * r_target)
        analytical_field = self.A0 * analytical_amplitude * np.sin(self.omega * t - self.kz * self.z)
        
        # 新增寻峰功能
        max_indices = argrelextrema(electric_field_z, np.greater, order=5)[0]
        
        # 计算波长
        if len(max_indices) >= 2:
            peak_positions = self.z[max_indices]
            wavelength = np.mean(np.diff(peak_positions))
            plt.text(0.05, 0.9, f'Wavelength: {wavelength:.4f} m', transform=plt.gca().transAxes, fontsize=12, color='blue')
        else:
            wavelength = None
        
        # 确保坐标轴范围固定
        plt.xlim(self.z[0], self.z[-1])
        plt.ylim(np.min(electric_field_z) * 1.1, np.max(electric_field_z) * 1.1)
        
        plt.figure(figsize=(10, 6))
        plt.plot(self.z, electric_field_z, label='numerical solution (Eθ)')
        plt.plot(self.z, analytical_field, label='analytical solution (Eθ)', linestyle='--')
        plt.plot(self.z, current_z, label='current density (Jθ)', linestyle='-.')
        plt.xlabel('z (m)')
        plt.ylabel('Field/Current')
        
        # 添加等离子体分界线
        plasma_boundary = self.length / 4
        plt.axvline(x=plasma_boundary, color='red', linestyle='--', linewidth=2, label='Plasma Boundary')
        plt.legend(loc='upper right')
        
        # 修改标题添加波长信息
        time_cycles = step * self.dt * self.frequency
        if wavelength:
            plt.title(f'Eθ and Jθ at r={r_target:.2f} m (t = {time_cycles:.2f} periods, λ = {wavelength:.4f} m)')
        else:
            plt.title(f'Eθ and Jθ at r={r_target:.2f} m (t = {time_cycles:.2f} periods)')
        
        # 添加极值标记（可选）
        if max_indices.size > 0:
            plt.scatter(self.z[max_indices], electric_field_z[max_indices], color='red', label='peaks')
        
        # 添加图例并将其位置设置为右下角
        plt.legend(loc='lower right')
        
        # 创建径向截面图的专用目录
        radial_dir = os.path.join(self.output_dir, 'radial_plots')
        os.makedirs(radial_dir, exist_ok=True)
        
        filename = f'field_{step:05d}_radial.png'
        plt.savefig(os.path.join(radial_dir, filename), dpi=100, bbox_inches='tight')
        plt.close()

    def create_radial_animation(self):
        """创建径向截面图动画"""
        print("Creating radial slice animation...")
        radial_files = sorted(glob.glob(os.path.join(self.output_dir, 'radial_plots', 'field_*.png')))
        
        if not radial_files:
            print("警告: 未找到径向截面图文件，无法生成动画")
            return
        
        # 读取所有径向截面图帧
        frames = [imageio.imread(file) for file in radial_files]
        
        # 保存动画
        output_path = os.path.join(self.output_dir, 'radial_animation.gif')
        try:
            imageio.mimsave(output_path, frames, duration=0.2)  # 每帧持续0.2秒
            print(f"径向截面图动画已保存到 {output_path}")
        except ValueError as e:
            print(f"保存GIF时出错: {e}")
            print("尝试使用替代方法...")
            
            # 使用PIL直接保存GIF
            from PIL import Image
            pil_frames = [Image.fromarray(frame) for frame in frames]
            pil_frames[0].save(
                output_path,
                save_all=True,
                append_images=pil_frames[1:],
                duration=200,  # 毫秒
                loop=0  # 0表示无限循环
            )
            print(f"使用PIL保存径向截面图动画到 {output_path}")

def main():
    try:
        frequency = 2.45e9  # 假设与C++中一致的频率值
        visualizer = WaveguideVisualizer(frequency)  # 传入频率参数

        # 查找场数据文件
        field_files = sorted(glob.glob('data/field_*.csv'))
        current_files = sorted(glob.glob('data/current_*.csv'))

        # 新增调试信息
        print(f"找到 {len(field_files)} 个场数据文件")
        print(f"找到 {len(current_files)} 个电流数据文件")

        if not field_files or not current_files:
            print("错误: 未找到场数据文件或电流数据文件")
            print("请检查以下内容：")
            print("- 确保当前工作目录为脚本所在目录")
            print("- 确保C++程序已正确运行并生成数据文件")
            print("- 确保数据文件命名符合 'field_*.csv' 和 'current_*.csv' 格式")
            return
        
        # 生成所有时间步的径向截面图
        print("生成所有时间步的径向截面图...")
        for field_file, current_file in zip(field_files, current_files):
            step = int(field_file.split('_')[1].split('.')[0])
            try:
                field = visualizer.load_field(field_file)
                current = visualizer.load_field(current_file)
                visualizer.plot_radial_slice(field, current, step)
            except Exception as e:
                print(f"处理文件 {field_file} 或 {current_file} 时出错: {e}")
        
        # 生成径向截面图动画
        print("准备生成径向截面图动画...")
        visualizer.create_radial_animation()

        # 生成动画
        print("准备生成动画序列...")
        visualizer.create_animation()
        
        print(f"\n可视化完成! 结果保存在 {visualizer.output_dir} 目录。")
            
    except FileNotFoundError as e:
        print(f"文件未找到错误: {e}")
    except ValueError as e:
        print(f"数值错误: {e}")
    except Exception as e:
        print(f"发生未预期的错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()